import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from sklearn.metrics import mean_squared_error, mean_absolute_error

# 设置字体为SimHei以显示中文
plt.rcParams['font.sans-serif'] = ['SimHei']
# 允许显示负号
plt.rcParams['axes.unicode_minus'] = False

# 创建保存图片的文件夹
output_dir = 'figures'
os.makedirs(output_dir, exist_ok=True)

# 读取Excel文件
file_path = 'pred/Ridge模型电力数据对比.xlsx'
sheet_names = ['Ridge']


# 定义计算评估指标的函数
def calculate_metrics(actual, predicted):
    # 确保输入是numpy数组
    actual = np.array(actual)
    predicted = np.array(predicted)

    # 过滤掉NaN值
    mask = ~np.isnan(predicted)
    actual = actual[mask]
    predicted = predicted[mask]

    if len(actual) == 0 or len(predicted) == 0:
        return np.nan, np.nan, np.nan

    rmse = np.sqrt(mean_squared_error(actual, predicted))
    mae = mean_absolute_error(actual, predicted)

    acc_li = []
    len_actual_pred = min(len(actual), len(predicted))
    for i in range(1, len_actual_pred):
        r = actual[i] - actual[i-1]
        p = predicted[i] - actual[i-1]
        if r * p < 0:
            acc_li.append(0)
        else:
            acc_li.append(1)
    
    return rmse, mae, np.mean(acc_li) if acc_li else np.nan


# 处理每个工作表
for sheet_name in sheet_names:
    print(f"\n{'=' * 50}")
    print(f"处理工作表: {sheet_name}")
    print(f"{'=' * 50}")

    # 读取数据
    df = pd.read_excel(file_path, sheet_name=sheet_name)

    # 清理列名（去除可能的空格）
    df.columns = df.columns.str.strip()

    # 确保日期列为datetime类型
    df['Date'] = pd.to_datetime(df['Date'])
    df.set_index('Date', inplace=True)

    # 分离训练集和测试集
    train_data = df[df['Set'] == 'Train']
    test_data = df[df['Set'] == 'Test']

    # 获取真实值和预测值
    y_train = train_data['True_Value']
    y_test = test_data['True_Value']

    # 原始预测值
    y_train_pred = train_data['Predicted_Value']
    y_test_pred = test_data['Predicted_Value']

    # 修正后的预测值（训练集与原始预测值相同，测试集使用Changed列）
    y_train_pred_corrected = train_data['Predicted_Value']  # 训练集不变
    y_test_pred_corrected = test_data['Changed']  # 测试集使用修正值

    # 计算原始预测值的评估指标
    train_rmse, train_mae, train_da = calculate_metrics(y_train, y_train_pred)
    test_rmse, test_mae, test_da = calculate_metrics(y_test, y_test_pred)

    # 计算修正后预测值的评估指标
    train_rmse_corr, train_mae_corr, train_da_corr = calculate_metrics(y_train, y_train_pred_corrected)
    test_rmse_corr, test_mae_corr, test_da_corr = calculate_metrics(y_test, y_test_pred_corrected)

    # 打印原始预测值的结果
    print(f"\n=== {sheet_name} - 原始预测值性能 ===")
    print(f"训练集 RMSE: {train_rmse:.6f}")
    print(f"训练集 MAE: {train_mae:.6f}")
    print(f"训练集 方向准确率: {train_da:.6%}" if not np.isnan(train_da) else "训练集 方向准确率: N/A")
    print(f"测试集 RMSE: {test_rmse:.6f}")
    print(f"测试集 MAE: {test_mae:.6f}")
    print(f"测试集 方向准确率: {test_da:.6%}" if not np.isnan(test_da) else "测试集 方向准确率: N/A")

    # 打印修正后预测值的结果
    print(f"\n=== {sheet_name} - 修正后预测值性能 ===")
    print(f"训练集 RMSE: {train_rmse_corr:.6f}")
    print(f"训练集 MAE: {train_mae_corr:.6f}")
    print(f"训练集 方向准确率: {train_da_corr:.6%}" if not np.isnan(train_da_corr) else "训练集 方向准确率: N/A")
    print(f"测试集 RMSE: {test_rmse_corr:.6f}")
    print(f"测试集 MAE: {test_mae_corr:.6f}")
    print(f"测试集 方向准确率: {test_da_corr:.6%}" if not np.isnan(test_da_corr) else "测试集 方向准确率: N/A")

    # 打印测试集真实值与原始预测值
    print(f"\n=== {sheet_name} - 测试集真实值与原始预测值 ===")
    comparison_df = pd.DataFrame({
        'True_Value': y_test.values,
        'Predicted_Value': y_test_pred.values
    }, index=y_test.index)
    print(comparison_df)

    # 打印测试集真实值与修正后预测值
    print(f"\n=== {sheet_name} - 测试集真实值与修正后预测值 ===")
    comparison_df_corrected = pd.DataFrame({
        'True_Value': y_test.values,
        'Corrected_Predicted_Value': y_test_pred_corrected.values
    }, index=y_test.index)
    print(comparison_df_corrected)

    # ========== 绘制并保存原始预测值的图表 ==========
    plt.figure(figsize=(14, 8))
    plt.plot(y_train.index, y_train.values, label='训练集真实值', color='blue', linewidth=2, marker='o', markersize=4)
    plt.plot(y_test.index, y_test.values, label='测试集真实值', color='green', linewidth=2, marker='s', markersize=4)
    plt.plot(y_train.index, y_train_pred.values, label='训练集修正预测值', color='orange', linestyle='--', linewidth=2,
             marker='^', markersize=4)
    plt.plot(y_test.index, y_test_pred.values, label='测试集修正预测值', color='red', linestyle='--', linewidth=2, marker='d',
             markersize=4)

    title_text = f'{sheet_name} 模型预测 GDP: 不变价: 当季同比\n测试集 RMSE: {test_rmse:.4f}, MAE: {test_mae:.4f}'
    if not np.isnan(test_da):
        title_text += f', 方向准确率: {test_da:.2%}'
    plt.title(title_text)
    plt.xlabel('时间')
    plt.ylabel('GDP: 不变价: 当季同比 (%)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xticks(rotation=45)
    plt.tight_layout()

    # 保存图片
    save_path = os.path.join(output_dir, f'{sheet_name}_原始预测.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"图表已保存至: {save_path}")
    plt.show()

    # ========== 绘制并保存修正后预测值的图表 ==========
    plt.figure(figsize=(14, 8))
    plt.plot(y_train.index, y_train.values, label='训练集真实值', color='blue', linewidth=2, marker='o', markersize=4)
    plt.plot(y_test.index, y_test.values, label='测试集真实值', color='green', linewidth=2, marker='s', markersize=4)
    plt.plot(y_train.index, y_train_pred_corrected.values, label='训练集修正预测值', color='orange', linestyle='--',
             linewidth=2, marker='^', markersize=4)
    plt.plot(y_test.index, y_test_pred_corrected.values, label='测试集修正预测值', color='red', linestyle='--', linewidth=2,
             marker='d', markersize=4)

    title_text = f'{sheet_name} 模型预测 GDP (修正后): 不变价: 当季同比\n测试集 RMSE: {test_rmse_corr:.4f}, MAE: {test_mae_corr:.4f}'
    if not np.isnan(test_da_corr):
        title_text += f', 方向准确率: {test_da_corr:.2%}'
    plt.title(title_text)
    plt.xlabel('时间')
    plt.ylabel('GDP: 不变价: 当季同比 (%)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xticks(rotation=45)
    plt.tight_layout()

    # 保存图片
    save_path = os.path.join(output_dir, f'{sheet_name}_修正后预测.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"图表已保存至: {save_path}")
    plt.show()